from importlib.metadata import entry_points
from typing import Set
import idaapi
from example.ida.utils import guess_machine,expr2colorstr
from future.utils import viewitems, viewvalues
from miasm.arch.x86.regs import *
from miasm.expression.simplifications import expr_simp
from miasm.core.asmblock import AsmCFG,AsmBlock,AsmConstraint
from miasm.ir.translators.translator import Translator
from miasm.ir.ir import IRCFG
from miasm.core.asmblock import AsmCFG

def color_irblock(irblock, lifter):
    out = []
    lbl = idaapi.COLSTR("%s:" % lifter.loc_db.pretty_str(irblock.loc_key), idaapi.SCOLOR_INSN)
    out.append(lbl)
    for assignblk in irblock:
        for dst, src in sorted(viewitems(assignblk)):
            dst_f = expr2colorstr(dst, loc_db=lifter.loc_db)
            src_f = expr2colorstr(src, loc_db=lifter.loc_db)
            line = idaapi.COLSTR("%s = %s" % (dst_f, src_f), idaapi.SCOLOR_INSN)
            out.append('    %s' % line)
        out.append("")
    out.pop()
    return "\n".join(out)

class GraphMiasmIR(idaapi.GraphViewer):

    def __init__(self, ircfg, title, result):
        idaapi.GraphViewer.__init__(self, title)
        self.ircfg = ircfg
        self.result = result
        self.names = {}

    def OnRefresh(self):
        self.Clear()
        addr_id = {}
        for (loc_key, irblock) in viewitems(self.ircfg.blocks):
            id_irblock = self.AddNode(color_irblock(irblock, self.ircfg))
            addr_id[loc_key] = id_irblock

        for (loc_key, irblock) in viewitems(self.ircfg.blocks):
            if not irblock:
                continue
            all_dst = self.ircfg.dst_trackback(irblock)
            for dst in all_dst:
                if not dst.is_loc():
                    continue
                if not dst.loc_key in self.ircfg.blocks:
                    continue
                node1 = addr_id[loc_key]
                node2 = addr_id[dst.loc_key]
                self.AddEdge(node1, node2)
        return True

    def OnGetText(self, node_id):
        return str(self[node_id])

    def OnSelect(self, _):
        return True

    def OnClick(self, _):
        return True

    def Show(self):
        if not idaapi.GraphViewer.Show(self):
            return False
        return True
def color_asmblock(asmblock:AsmBlock):
    out = []
    lbl = idaapi.COLSTR("%s:" % asmblock.loc_db.pretty_str(asmblock.loc_key), idaapi.SCOLOR_INSN)
    out.append(asmblock.loc_db.pretty_str(asmblock.loc_key))
    for line in asmblock.lines:
        line_str = idaapi.COLSTR("%s" % line, idaapi.SCOLOR_INSN)
        out.append("    %s" % line_str)
    # out.pop()
    return "\n".join(out)
class GraphAsmCfg(idaapi.GraphViewer):
    def __init__(self, asmcfg:AsmCFG, title, result):
        idaapi.GraphViewer.__init__(self, title)
        self.asmcfg = asmcfg
        self.result = result
        self.names = {}

    def OnRefresh(self):
        self.Clear()
        addr_id = {}
        for asmblock in self.asmcfg.blocks:
            id_asmblock=self.AddNode(color_asmblock(asmblock))
            addr_id[asmblock.loc_key]=id_asmblock
        for asmblock in self.asmcfg.blocks:
            loc_key=asmblock.loc_key 
            if not asmblock:
                continue
            for bto in asmblock.bto:
                node1=addr_id[loc_key]
                node2=addr_id[bto.loc_key]
                self.AddEdge(node1,node2)
        return True
    def OnGetText(self, node_id):
        return str(self[node_id])

    def OnSelect(self, _):
        return True

    def OnClick(self, _):
        return True

    def Show(self):
        if not idaapi.GraphViewer.Show(self):
            return False
        return True
    
def draw_graph(cfg,name):
    if isinstance(cfg,AsmCFG):
        graph=GraphAsmCfg(cfg,name,None)
        graph.Show()
    elif isinstance(cfg,IRCFG):
        graph=GraphMiasmIR(cfg,name,None)
        graph.Show()
    else:
        print("[x]unknow cfg type")